# -*- coding:utf-8 -*-
from typing import Optional
from contextlib import contextmanager

from sqlmodel import create_engine
from sqlmodel import Session as SQLSession
from redis_om import get_redis_connection

from .exceptions import ConfigurationError
from .models import StorageMode, OrmModel

class HybridEngine:
    def __init__(
            self,
            sql_url: Optional[str] = None,
            redis_url: Optional[str] = "redis://localhost:6379/0"
    ):
        self.sql_engine = create_engine(sql_url) if sql_url else None
        self.redis = get_redis_connection(url=redis_url) if redis_url else None

    @contextmanager
    def session(self):
        """上下文管理混合会话"""
        session = HybridSession(self)
        try:
            # session = HybridSession(self)
            yield session
            session.commit()
        except Exception as e:
            session.rollback()
            raise e
        finally:
            session.close()


class HybridSession(SQLSession):
    def __init__(self, engine: HybridEngine):
        super().__init__(engine.sql_engine if engine.sql_engine else None)
        self.engine = engine
        self._redis_dirty_keys = set()  # 记录需回滚的 Redis 键

    def _validate_storage_mode(self, instance: OrmModel) -> None:
        """验证存储模式与引擎配置是否兼容"""
        mode = instance.__storage_mode__
        if mode == StorageMode.SQL and not self.engine.sql_engine:
            raise ConfigurationError("SQL 存储模式需要配置 SQL 数据库连接")
        if mode == StorageMode.REDIS and not self.engine.redis:
            raise ConfigurationError("Redis 存储模式需要配置 Redis 连接")

    def commit(self) -> None:
        """提交事务"""
        try:
            if self.engine.sql_engine:
                super().commit()
            # Redis 无显式事务，直接持久化
        except Exception as e:
            self.rollback()
            raise e

    def rollback(self) -> None:
        """回滚事务"""
        if self.engine.sql_engine:
            super().rollback()
        # Redis 数据回滚（删除脏数据）
        if self.engine.redis and self._redis_dirty_keys:
            self.engine.redis.delete(*self._redis_dirty_keys)
            self._redis_dirty_keys.clear()

    def close(self) -> None:
        """关闭会话"""
        if self.engine.sql_engine:
            super().close()
        # Redis 连接由连接池管理，无需关闭

    def save(self, instance: OrmModel) -> None:
        """保存数据（新增/更新）"""
        self._validate_storage_mode(instance)
        key = instance.redis_key()

        # SQL 操作
        if instance.__storage_mode__ in [StorageMode.SQL, StorageMode.HYBRID]:
            if not instance.id:  # 处理自增 ID
                self.add(instance)
                self.flush()  # 强制生成 ID
            else:
                self.merge(instance)

        # Redis 操作
        if instance.__storage_mode__ in [StorageMode.REDIS, StorageMode.HYBRID]:
            self.engine.redis.hset(key, mapping=instance.model_dump())
            if instance.__redis_ttl__:
                self.engine.redis.expire(key, instance.__redis_ttl__)
            self._redis_dirty_keys.add(key)  # 标记为可能需回滚